import pandas as pd
import requests
import time
import re
from urllib.parse import quote

def replace_greek_letters(text):
    """Replace Greek letters with their English word equivalents"""
    if pd.isna(text):
        return text
        
    replacements = {
        'α': 'alpha',
        'β': 'beta', 
        'γ': 'gamma',
        'δ': 'delta',
        'ε': 'epsilon',
        'ζ': 'zeta',
        'η': 'eta',
        'θ': 'theta',
        'ι': 'iota',
        'κ': 'kappa',
        'λ': 'lambda',
        'μ': 'mu',
        'ν': 'nu',
        'ξ': 'xi',
        'ο': 'omicron',
        'π': 'pi',
        'ρ': 'rho',
        'σ': 'sigma',
        'τ': 'tau',
        'υ': 'upsilon',
        'φ': 'phi',
        'χ': 'chi',
        'ψ': 'psi',
        'ω': 'omega',
        'Α': 'Alpha',
        'Β': 'Beta',
        'Γ': 'Gamma',
        'Δ': 'Delta',
        'Ε': 'Epsilon',
        'Ζ': 'Zeta',
        'Η': 'Eta',
        'Θ': 'Theta',
        'Ι': 'Iota',
        'Κ': 'Kappa',
        'Λ': 'Lambda',
        'Μ': 'Mu',
        'Ν': 'Nu',
        'Ξ': 'Xi',
        'Ο': 'Omicron',
        'Π': 'Pi',
        'Ρ': 'Rho',
        'Σ': 'Sigma',
        'Τ': 'Tau',
        'Υ': 'Upsilon',
        'Φ': 'Phi',
        'Χ': 'Chi',
        'Ψ': 'Psi',
        'Ω': 'Omega'
    }
    
    result = text
    for greek, english in replacements.items():
        result = result.replace(greek, english)
    
    return result

def extract_cas_from_title(title):
    """Extract CAS number from a title string"""
    # Pattern for CAS number: numbers-numbers-number
    cas_pattern = r'\b\d+-\d+-\d+\b'
    match = re.search(cas_pattern, title)
    return match.group(0) if match else ''

def search_pubchem(compound_name):
    """
    Search PubChem for a compound and return its CID, CAS, MF (Molecular Formula), and SMILES.
    """
    try:
        if not compound_name or pd.isna(compound_name):
            print(f"Skipping empty compound name")
            return {'CID': '', 'SMILES': '', 'CAS': '', 'MF': ''}
            
        # Replace Greek letters with English words for searching
        search_name = replace_greek_letters(compound_name)
        print(f"Searching for: {search_name} (original: {compound_name})")
        
        # URL encode the compound name
        encoded_name = quote(search_name)
        
        # First, search for the compound to get its CID
        search_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{encoded_name}/cids/JSON"
        response = requests.get(search_url)
        
        if response.status_code == 200 and 'IdentifierList' in response.json():
            cid = str(response.json()['IdentifierList']['CID'][0])  # Get the best match
            
            # Get the synonyms to extract CAS number
            synonyms_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/synonyms/JSON"
            synonyms_response = requests.get(synonyms_url)
            
            cas = ''
            if synonyms_response.status_code == 200:
                synonyms = synonyms_response.json()['InformationList']['Information'][0]['Synonym']
                # Look for CAS number in synonyms
                for synonym in synonyms:
                    cas_match = extract_cas_from_title(synonym)
                    if cas_match:
                        cas = cas_match
                        break
            
            # Get SMILES using the CID
            property_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/CanonicalSMILES,MolecularFormula/JSON"
            prop_response = requests.get(property_url)
            
            if prop_response.status_code == 200:
                properties = prop_response.json()['PropertyTable']['Properties'][0]
                smiles = properties['CanonicalSMILES']
                molecular_formula = properties['MolecularFormula']
                
                return {
                    'CID': cid,
                    'SMILES': smiles,
                    'CAS': cas,
                    'MF': molecular_formula
                }
        
        print(f"No data found for {compound_name}")
        return {'CID': '', 'SMILES': '', 'CAS': '', 'MF': ''}
    
    except Exception as e:
        print(f"Error processing {compound_name}: {str(e)}")
        return {'CID': '', 'SMILES': '', 'CAS': '', 'MF': ''}

def update_excel_with_pubchem_data(input_file, output_file):
    """
    Update Excel file with PubChem data using the specified column mappings.
    """
    # Read the Excel file
    print(f"Reading input file: {input_file}")
    df = pd.read_excel(input_file)
    
    # Check if the required columns exist
    required_column = "英文名"
    if required_column not in df.columns:
        print(f"Error: Required column '{required_column}' not found in the Excel file.")
        print(f"Available columns: {', '.join(df.columns)}")
        return
    
    # Initialize columns if they don't exist
    target_columns = {
        'CID': 'CID号',
        'CAS': 'CAS号',
        'MF': '分子式',
        'SMILES': 'SMILES'
    }
    
    for _, col_name in target_columns.items():
        if col_name not in df.columns:
            df[col_name] = ''
    
    # Store original compound names to preserve Greek letters
    original_names = {}
    for idx, row in df.iterrows():
        name = str(row[required_column]).strip()
        if not name or name == 'nan':
            continue
        original_names[replace_greek_letters(name)] = name
    
    # Get unique compound names from 英文名 column
    # Remove NaN values and convert to list
    unique_compounds = df[required_column].dropna().unique().tolist()
    print(f"Found {len(unique_compounds)} unique compounds out of {len(df)} total entries")
    
    # Create a dictionary to store results for each unique compound
    compound_data = {}
    
    # Process unique compounds
    total = len(unique_compounds)
    for idx, compound_name in enumerate(unique_compounds):
        compound_name = str(compound_name).strip()
        if not compound_name:
            continue
            
        print(f"Processing {idx+1}/{total}: {compound_name}")
        
        # Search PubChem
        result = search_pubchem(compound_name)
        
        # Store results using the original compound name (with Greek letters)
        compound_data[compound_name] = result
        
        # Print results for debugging
        print(f"Found data: CID={result['CID']}, CAS={result['CAS']}, MF={result['MF']}, SMILES={result['SMILES']}")
        
        # Add a small delay to avoid overwhelming the PubChem API
        time.sleep(0.3)
    
    # Fill in the target columns using the cached results
    print("Filling in data for all entries...")
    for idx, row in df.iterrows():
        name = str(row[required_column]).strip()
        if not name or name == 'nan':
            continue
            
        if name in compound_data:
            data = compound_data[name]
            df.at[idx, target_columns['CID']] = data['CID']
            df.at[idx, target_columns['CAS']] = data['CAS']
            df.at[idx, target_columns['MF']] = data['MF']
            df.at[idx, target_columns['SMILES']] = data['SMILES']
    
    # Save the updated dataframe to the output file
    df.to_excel(output_file, index=False)
    print(f"Processing complete. Results saved to {output_file}")
    
    # Print some statistics
    successful_lookups = sum(1 for data in compound_data.values() if data['CID'])
    print(f"\nStatistics:")
    print(f"Total unique compounds processed: {total}")
    print(f"Successful lookups: {successful_lookups}")
    print(f"Failed lookups: {total - successful_lookups}")


# Usage example
if __name__ == "__main__":
    input_file = "./API/蔓荆子0320_matched.xlsx"  
    output_file = "./API/蔓荆子0320_matched_PC.xlsx"  
    
    update_excel_with_pubchem_data(input_file, output_file)